package com.lifeonwalden.copula;

import com.lifeonwalden.copula.bean.OptimizeBounds;
import com.lifeonwalden.copula.bean.OptimizeOptions;
import com.lifeonwalden.copula.constant.ArchimedeanCopulaFamily;
import com.lifeonwalden.copula.constant.ArchimedeanCopulaMethod;
import com.lifeonwalden.copula.constant.CopulaConstants;
import com.lifeonwalden.copula.constant.OptimizeMinimizeScalarMethod;
import com.lifeonwalden.copula.func.ArchimedeanDerivative;
import com.lifeonwalden.copula.func.ArchimedeanGenerators;
import com.lifeonwalden.copula.util.ArrayUtil;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;

public class ArchimedeanCopula extends Copula {
    private ArchimedeanCopulaFamily family;
    private BigDecimal parameter;

    public ArchimedeanCopula(ArchimedeanCopulaFamily family) {
        switch (family) {
            case AMH: {
                this.parameter = BigDecimal.valueOf(0.5);
                break;
            }
            default: {
                this.parameter = BigDecimal.valueOf(1.5);
            }
        }
        this.family = family;
    }

    public BigDecimal generator(BigDecimal x) {
        switch (this.family) {
            case AMH: {
                return ArchimedeanGenerators.aliMikhailHaqGenerator(x, this.parameter);
            }
            case JOE: {
                return ArchimedeanGenerators.joeGenerator(x, this.parameter);
            }
            case FRANK: {
                return ArchimedeanGenerators.frankGenerator(x, this.parameter);
            }
            case GUMBEL: {
                return ArchimedeanGenerators.gumbelGenerator(x, this.parameter);
            }
            default: {
                return ArchimedeanGenerators.claytonGenerator(x, this.parameter);
            }
        }
    }

    public BigDecimal inverseGenerator(BigDecimal x) {
        switch (this.family) {
            case AMH: {
                return ArchimedeanGenerators.aliMikhailHaqGeneratorInvert(x, this.parameter);
            }
            case JOE: {
                return ArchimedeanGenerators.joeGeneratorInvert(x, this.parameter);
            }
            case FRANK: {
                return ArchimedeanGenerators.frankGeneratorInvert(x, this.parameter);
            }
            case GUMBEL: {
                return ArchimedeanGenerators.gumbelGeneratorInvert(x, this.parameter);
            }
            default: {
                return ArchimedeanGenerators.claytonGeneratorInvert(x, this.parameter);
            }
        }
    }

    public BigDecimal getParameter() {
        return parameter;
    }

    public ArchimedeanCopula setParameter(BigDecimal theta) {
        this.parameter = theta;

        return this;
    }

    public ArchimedeanCopulaFamily getFamily() {
        return family;
    }

    public void fit(List<BigDecimal>[] sampleData, ArchimedeanCopulaMethod method, Optional<OptimizeBounds> thetaBounds) {
        switch (method) {
            // Canonical Maximum Likelihood Estimation
            case CMLE: {
                // Pseudo-observations from real data
                BigDecimal[][] pobs = new BigDecimal[sampleData.length][sampleData[0].size()];
                for (int i = 0; i < sampleData.length; i++) {
                    int[] ranks = ArrayUtil.fetchIndexOfAscOrder(ArrayUtil.fetchIndexOfAscOrder(sampleData[i]));
                    for (int k = 0; k < ranks.length; k++) {
                        pobs[i][k] = BigDecimal.ONE.add(BigDecimal.valueOf(ranks[k]))
                                .divide(BigDecimal.ONE.add(BigDecimal.valueOf(ranks.length)), 16, RoundingMode.DOWN);
                    }
                }

                Function<BigDecimal, BigDecimal> logLikelihoodFunc = theta -> {
                    double[] paramObs = new double[pobs[0].length];
                    for (int i = 0; i < pobs[0].length; i++) {
                        paramObs[i] = this.pdfParam(new BigDecimal[]{pobs[0][i], pobs[1][i]}, theta).doubleValue();
                    }

                    double sum = 0;
                    for (double paramOb : paramObs) {
                        sum += Math.log(paramOb);
                    }

                    return BigDecimal.valueOf(-1 * sum);
                };

                boolean isScalar = true;
                BigDecimal thetaStart = CopulaConstants.POINT_FIVE;
                OptimizeBounds bounds = new OptimizeBounds();
                if (thetaBounds.isPresent()) {
                    bounds = bounds.setUpper(thetaBounds.get().getUpper()).setLower(thetaBounds.get().getLower());
                } else {
                    switch (this.family) {
                        case AMH: {
                            bounds.setLower(CopulaConstants.MINUS_ONE).setUpper(BigDecimal.ONE.subtract(BigDecimal.valueOf(1 - 1e-6)));
                            isScalar = false;
                            thetaStart = CopulaConstants.POINT_FIVE;
                            break;
                        }
                        case CLAYTON: {
                            bounds.setLower(BigDecimal.ZERO).setUpper(BigDecimal.TEN);
                            break;
                        }
                        case GUMBEL:
                        case JOE: {
                            bounds.setLower(BigDecimal.ONE).setUpper(null);
                            isScalar = false;
                            thetaStart = BigDecimal.valueOf(1.5);
                            break;
                        }
                    }
                }

                this.parameter = Estimation
                        .cmle(logLikelihoodFunc, thetaStart, Optional.of(bounds), new OptimizeOptions().setScalar(isScalar).setMinimizeScalarMethod(OptimizeMinimizeScalarMethod.BOUNDED))
                        .getX();

                break;
            }
            default: {
                throw new RuntimeException("Not support yet");
            }
        }
    }

    @Override
    public BigDecimal cdf(double[] data) {
        return null;
    }

    @Override
    public BigDecimal pdf(double[] data) {
        return null;
    }

    /**
     * Returns the PDF of the copula with the specified theta. Use this when you want to compute PDF with another parameter.
     *
     * @param quantiles Quantiles
     * @param theta     The custom parameter.
     * @return
     */
    public BigDecimal pdfParam(BigDecimal[] quantiles, BigDecimal theta) {
        // prod is the product of the derivatives of the generator for each variable
        BigDecimal prod = BigDecimal.ONE;

        // The sum of generators that will be computed on the invert derivative
        BigDecimal sumInvert = BigDecimal.ZERO;

        // the n-th derivative of the invert
        BigDecimal invertNDerivative;
        if ((this.family == ArchimedeanCopulaFamily.FRANK || this.family == ArchimedeanCopulaFamily.CLAYTON) && theta.compareTo(BigDecimal.ZERO) == 0) {
            theta = BigDecimal.valueOf(1e-8);
        }

        switch (this.family) {
            case CLAYTON: {
                // We compute product and sum
                for (BigDecimal quantile : quantiles) {
                    prod = prod.multiply(BigDecimal.valueOf(-1 * Math.pow(quantile.doubleValue(), -theta.doubleValue() - 1)));
                    sumInvert = sumInvert.add(ArchimedeanGenerators.claytonGenerator(quantile, theta));
                }

                invertNDerivative = ArchimedeanDerivative.claytonInvertNDerivative(sumInvert, theta, quantiles.length);
                break;
            }
            default: {
                throw new RuntimeException("Not support yet");
            }
        }

        return prod.multiply(invertNDerivative);
    }

    @Override
    public BigDecimal concentrationDown(double x) {
        return null;
    }

    @Override
    public BigDecimal concentrationUp(double x) {
        return null;
    }
}
